/*
 * Copyright [2022] [https://www.xiaonuo.vip]
 *
 * Snowy采用APACHE LICENSE 2.0开源协议，您在使用过程中，需要注意以下几点：
 *
 * 1.请不要删除和修改根目录下的LICENSE文件。
 * 2.请不要删除和修改Snowy源码头部的版权声明。
 * 3.本项目代码可免费商业使用，商业使用请保留源码和相关描述文件的项目出处，作者声明等。
 * 4.分发源码时候，请注明软件出处 https://www.xiaonuo.vip
 * 5.不可二次分发开源参与同类竞品，如有想法可联系团队xiaonuobase@qq.com商议合作。
 * 6.若您的项目无法满足以上几点，需要更多功能代码，获取Snowy商业授权许可，请在官网购买授权，地址为 https://www.xiaonuo.vip
 */
package org.wenshu.ai.modular.traindata.service.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollStreamUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.wenshu.ai.modular.chat.provider.core.AiModelService;
import org.wenshu.common.enums.CommonSortOrderEnum;
import org.wenshu.common.exception.CommonException;
import org.wenshu.common.page.CommonPageRequest;
import org.wenshu.ai.modular.traindata.entity.TrainData;
import org.wenshu.ai.modular.traindata.mapper.TrainDataMapper;
import org.wenshu.ai.modular.traindata.param.TrainDataAddParam;
import org.wenshu.ai.modular.traindata.param.TrainDataEditParam;
import org.wenshu.ai.modular.traindata.param.TrainDataIdParam;
import org.wenshu.ai.modular.traindata.param.TrainDataPageParam;
import org.wenshu.ai.modular.traindata.service.TrainDataService;

import java.util.List;
import java.util.UUID;

/**
 * 训练数据表Service接口实现类
 *
 * @author poker
 * @date  2025/03/12 14:08
 **/
@Service
public class TrainDataServiceImpl extends ServiceImpl<TrainDataMapper, TrainData> implements TrainDataService {

    @Resource
    private AiModelService aiModelService;

    @Override
    public Page<TrainData> page(TrainDataPageParam trainDataPageParam) {
        QueryWrapper<TrainData> queryWrapper = new QueryWrapper<TrainData>().checkSqlInjection();
        if(ObjectUtil.isNotEmpty(trainDataPageParam.getQuestion())) {
            queryWrapper.lambda().like(TrainData::getQuestion, trainDataPageParam.getQuestion());
        }
        if(ObjectUtil.isNotEmpty(trainDataPageParam.getSqlText())) {
            queryWrapper.lambda().like(TrainData::getSqlText, trainDataPageParam.getSqlText());
        }
        if(ObjectUtil.isNotEmpty(trainDataPageParam.getDdlText())) {
            queryWrapper.lambda().like(TrainData::getDdlText, trainDataPageParam.getDdlText());
        }
        if(ObjectUtil.isNotEmpty(trainDataPageParam.getDataType())) {
            queryWrapper.lambda().eq(TrainData::getDataType, trainDataPageParam.getDataType());
        }
        if(ObjectUtil.isNotEmpty(trainDataPageParam.getDocument())) {
            queryWrapper.lambda().like(TrainData::getDocument, trainDataPageParam.getDocument());
        }
        if(ObjectUtil.isNotEmpty(trainDataPageParam.getIsPositiveExample())) {
            queryWrapper.lambda().eq(TrainData::getIsPositiveExample, trainDataPageParam.getIsPositiveExample());
        }
        if(ObjectUtil.isAllNotEmpty(trainDataPageParam.getSortField(), trainDataPageParam.getSortOrder())) {
            CommonSortOrderEnum.validate(trainDataPageParam.getSortOrder());
            queryWrapper.orderBy(true, trainDataPageParam.getSortOrder().equals(CommonSortOrderEnum.ASC.getValue()),
                    StrUtil.toUnderlineCase(trainDataPageParam.getSortField()));
        } else {
            queryWrapper.lambda().orderByAsc(TrainData::getId);
        }
        return this.page(CommonPageRequest.defaultPage(), queryWrapper);
    }

    private void storeEmbedding(String dataType, String content, String vectorId) {
        Embedding embedding = aiModelService.getEmbeddingModel().embed(content).content();
        List<Embedding> embeddings = List.of(embedding);
        List<TextSegment> texts = List.of(TextSegment.from(content));
        List<String> vectorIds = List.of(vectorId);

        EmbeddingStore<TextSegment> store;
        if ("ddl".equals(dataType)) {
            store = aiModelService.getDdlEmbeddingStore();
        } else if ("sql".equals(dataType)) {
            store = aiModelService.getSqlEmbeddingStore();
        } else if ("doc".equals(dataType)) {
            store = aiModelService.getDocEmbeddingStore();
        } else {
            throw new CommonException("不支持的数据类型:" + dataType);
        }

        store.addAll(vectorIds, embeddings, texts);
    }

    @Transactional(rollbackFor = Exception.class)
    @Override
    public void add(TrainDataAddParam trainDataAddParam) {
        TrainData trainData = BeanUtil.toBean(trainDataAddParam, TrainData.class);
        // 生成向量ID
        trainData.setVectorId(UUID.randomUUID().toString());
        this.save(trainData);

        String content = null;
        if ("ddl".equals(trainData.getDataType())) {
            content = trainData.getDdlText();
        } else if ("sql".equals(trainData.getDataType())) {
            // 将问题向量化
            content = trainData.getQuestion();
        } else if ("doc".equals(trainData.getDataType())) {
            content = trainData.getDocument();
        }

        if (content != null) {
            storeEmbedding(trainData.getDataType(), content, trainData.getVectorId());
        }
    }

    @Transactional(rollbackFor = Exception.class)
    @Override
    public void edit(TrainDataEditParam trainDataEditParam) {
        TrainData trainData = this.queryEntity(trainDataEditParam.getId());
        // 如果向量ID为空，生成新的
        if (StrUtil.isEmpty(trainData.getVectorId())) {
            trainData.setVectorId(UUID.randomUUID().toString());
        }
        BeanUtil.copyProperties(trainDataEditParam, trainData);
        this.updateById(trainData);

        String content = null;
        if ("ddl".equals(trainData.getDataType())) {
            content = trainData.getDdlText();
        } else if ("sql".equals(trainData.getDataType())) {
            // 将问题向量化
            content = trainData.getQuestion();
        } else if ("doc".equals(trainData.getDataType())) {
            content = trainData.getDocument();
        }

        if (content != null) {
            storeEmbedding(trainData.getDataType(), content, trainData.getVectorId());
        }
    }

    @Transactional(rollbackFor = Exception.class)
    @Override
    public void delete(List<TrainDataIdParam> trainDataIdParamList) {
        // 获取要删除的数据
        List<TrainData> trainDataList = this.listByIds(CollStreamUtil.toList(trainDataIdParamList, TrainDataIdParam::getId));

        // 从向量数据库中删除
        for (TrainData trainData : trainDataList) {
            if (StrUtil.isNotEmpty(trainData.getVectorId())) {
                if ("ddl".equals(trainData.getDataType())) {
                    aiModelService.getDdlEmbeddingStore().remove(trainData.getVectorId());
                } else if ("sql".equals(trainData.getDataType())) {
                    aiModelService.getSqlEmbeddingStore().remove(trainData.getVectorId());
                } else if ("doc".equals(trainData.getDataType())) {
                    aiModelService.getDocEmbeddingStore().remove(trainData.getVectorId());
                }
            }
        }

        // 执行删除
        this.removeByIds(CollStreamUtil.toList(trainDataIdParamList, TrainDataIdParam::getId));
    }

    @Override
    public TrainData detail(TrainDataIdParam trainDataIdParam) {
        return this.queryEntity(trainDataIdParam.getId());
    }

    @Override
    public TrainData queryEntity(String id) {
        TrainData trainData = this.getById(id);
        if(ObjectUtil.isEmpty(trainData)) {
            throw new CommonException("训练数据表不存在，id值为：{}", id);
        }
        return trainData;
    }
}
